package org.apache.hadoop.hdfs.shortcircuit;

import java.io.BufferedOutputStream;
import java.io.Closeable;
import java.io.DataOutputStream;
import java.io.EOFException;
import java.io.FileInputStream;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map.Entry;
import java.util.TreeMap;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;

import org.apache.commons.lang.mutable.MutableBoolean;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.classification.InterfaceAudience;
import org.apache.hadoop.hdfs.ExtendedBlockId;
import org.apache.hadoop.hdfs.net.DomainPeer;
import org.apache.hadoop.hdfs.protocol.DatanodeInfo;
import org.apache.hadoop.hdfs.protocol.datatransfer.DataTransferProtocol;
import org.apache.hadoop.hdfs.protocol.datatransfer.Sender;
import org.apache.hadoop.hdfs.protocol.proto.DataTransferProtos.ShortCircuitShmResponseProto;
import org.apache.hadoop.hdfs.protocolPB.PBHelper;
import org.apache.hadoop.hdfs.server.datanode.ShortCircuitRegistry;
import org.apache.hadoop.hdfs.shortcircuit.ShortCircuitShm.ShmId;
import org.apache.hadoop.hdfs.shortcircuit.ShortCircuitShm.Slot;
import org.apache.hadoop.io.IOUtils;
import org.apache.hadoop.net.unix.DomainSocket;
import org.apache.hadoop.net.unix.DomainSocketWatcher;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;

/**
 * DfsClientShmManager管理着HDFS客户端侧的所有共享内存，
 *
 * ![](https://markdown-sjc.oss-cn-beijing.aliyuncs.com/img/20210618163127.png)
 *
 * 一个DFSClient可能需要维护与多个Datanode之间的共享内存，所以DfsClientShmManager定义了内部类EndpointShmManager来管理与一个Datanode之间的所有共享内存，每段共享内存又由一个DfsClientShm对象管理.
 *
 * DfsClientShmManager  定义了datanodes字段保存所有的EndPointShmManager对象，
 * 同时对外提供allocSlot()和freeSlot()两个方法分别用于从指定Datanode的共享内存中分配和释放一个槽位。
 * allocSlot()和freeSlot()方法首先都会从datanodes字段中提取出该Datanode对应的EndpointShmManager对象，然后调用EndpointShmManager的allocSlot()和freeSlot()方法执行获取和释放槽位的操作。
 *
 * DFSClient和一个Datanode之间可能有多段共享内存，其中有些共享内存可能已经没有可用槽位了，所以EndpointShmManager定义了两个队列：notFul队列存储还有槽位的共享内存；full队列存储没有槽位的共享内存。
 *
 * EndpointShmManager类还定义了如下几个方法。
 *     allocSlot()方法：用来在共享内存中申请一个槽位。
 *     freeSlot()方法：用于释放一个槽位。allocSlot()和freeSLot()这两个方法都是在ShortCircuitCache创建或者关闭一个ShortCircuitReplica时调用的。
 *     requestNewShm()方法：用于与Datanode协商创建一段新的共享内存。
 */
@InterfaceAudience.Private
public class DfsClientShmManager implements Closeable {
  private static final Log LOG = LogFactory.getLog(DfsClientShmManager.class);

  /**
   * Manages short-circuit memory segments that pertain to a given DataNode.
   */
  class EndpointShmManager {
    /**
     * The datanode we're managing.
     */
    private final DatanodeInfo datanode;

    /**
     * Shared memory segments which have no empty slots.
     *
     * Protected by the manager lock.
     */
    private final TreeMap<ShmId, DfsClientShm> full =
        new TreeMap<ShmId, DfsClientShm>();

    /**
     * Shared memory segments which have at least one empty slot.
     *
     * Protected by the manager lock.
     */
    private final TreeMap<ShmId, DfsClientShm> notFull =
        new TreeMap<ShmId, DfsClientShm>();

    /**
     * True if this datanode doesn't support short-circuit shared memory
     * segments.
     *
     * Protected by the manager lock.
     */
    private boolean disabled = false;

    /**
     * True if we're in the process of loading a shared memory segment from
     * this DataNode.
     *
     * Protected by the manager lock.
     */
    private boolean loading = false;

    EndpointShmManager (DatanodeInfo datanode) {
      this.datanode = datanode;
    }

    /**
     * Pull a slot out of a preexisting shared memory segment.
     *
     * Must be called with the manager lock held.
     *
     * @param blockId     The blockId to put inside the Slot object.
     *
     * @return            null if none of our shared memory segments contain a free slot; the slot object otherwise.
     */
    private Slot allocSlotFromExistingShm(ExtendedBlockId blockId) {
      if (notFull.isEmpty()) {
        return null;
      }
      Entry<ShmId, DfsClientShm> entry = notFull.firstEntry();
      DfsClientShm shm = entry.getValue();
      ShmId shmId = shm.getShmId();
      Slot slot = shm.allocAndRegisterSlot(blockId);
      if (shm.isFull()) {
        DfsClientShm removedShm = notFull.remove(shmId);
        Preconditions.checkState(removedShm == shm);
        full.put(shmId, shm);
      } else {
      }
      return slot;
    }

    /**
     * requestNewShm()方法用于在共享内存中没有可用槽位时，与Datanode协商构造一段新的共享内存。
     *
     * requestNewShm()方法会调用DataTransferProtocol.requestShortCircuitShm()方法向Datanode申请一段共享内存，
     * Datanode会通过domainSocket将共享内存文件的文件描述符发回EndpointShmManager。
     * EndpointShmManager会打开共享内存文件并执行内存映射操作，然后EndpointShmManager会构造DfsClientShm对象并返回。
     */
    private DfsClientShm requestNewShm(String clientName, DomainPeer peer) throws IOException {
      final DataOutputStream out = new DataOutputStream( new BufferedOutputStream(peer.getOutputStream()));
      // 调用DataTransferProtocol.requestShortCircuitShm() 方法向Datanode申请一段共享内存
      new Sender(out).requestShortCircuitShm(clientName);
      // 接收响应消息
      ShortCircuitShmResponseProto resp = ShortCircuitShmResponseProto.parseFrom( PBHelper.vintPrefixed(peer.getInputStream()));
      String error = resp.hasError() ? resp.getError() : "(unknown)";
      switch (resp.getStatus()) {
      case SUCCESS: // 如果成功
        DomainSocket sock = peer.getDomainSocket();
        byte buf[] = new byte[1];
        // datanode返回的共享内存段的文件描述符是通过一个文件输入流来描述的，
        // 所以此时是通过socket直接获取那个文件输入流的
        FileInputStream fis[] = new FileInputStream[1];
        // 从 domainSocket 接收共享内存文件的文件描述符
        if (sock.recvFileInputStreams(fis, buf, 0, buf.length) < 0) {
          throw new EOFException("got EOF while trying to transfer the " + "file descriptor for the shared memory segment.");
        }
        if (fis[0] == null) {
          throw new IOException("the datanode " + datanode + " failed to " + "pass a file descriptor for the shared memory segment.");
        }
        try {
          // 构造 DfsClientShm对象
          DfsClientShm shm = new DfsClientShm(PBHelper.convert(resp.getId()), fis[0], this, peer);
          return shm;
        } finally {
          IOUtils.cleanup(LOG,  fis[0]);
        }
      case ERROR_UNSUPPORTED:
        // The DataNode just does not support short-circuit shared memory access, and we should stop asking.
        LOG.info(this + ": datanode does not support short-circuit " + "shared memory access: " + error);
        disabled = true;
        return null;
      default:
        // The datanode experienced some kind of unexpected error when trying to create the short-circuit shared memory segment.
        LOG.warn(this + ": error requesting short-circuit shared memory " + "access: " + error);
        return null;
      }
    }

    /**
     * allocSlot()方法实现了分配槽位的操作。
     * 这个方法首先从notFull队列中的共享内存中申请一个槽位。如果当前EndpointShmManager管理的所有共享内存中已经没有槽位了，则调用requestNewShm()方法通过DataTransferProtocol与Datanode申请一段新的共享内存。
     * 申请完成后，构造DfsClientShm对象管理这段共享内存，同时触发DfsClientShm监控底层的domainSocket，并将这个新构造的DfsClientShm对象放入notFull队列中。
     */
    Slot allocSlot(DomainPeer peer, MutableBoolean usedPeer, String clientName, ExtendedBlockId blockId) throws IOException {
      while (true) {
        if (closed) {
          // DfsClientShmManager已经关闭了，返回null
          return null;
        }
        if (disabled) {
          // 如果当前EndpointShmManager对应的Datanode不支持短路读取，则返回null
          return null;
        }
        // 如果已经申请了共享内存，则从该共享内存中申请槽位
        // 一开始是没有共享内存段的，所以刚开始是获取不到slot的，必然申请一个新的共享内存
        // 或者是一个共享内存段的slot都用光了，此时也申请一个新的共享内存段
        Slot slot = allocSlotFromExistingShm(blockId); // TODO  此时就可以从刚申请的共享内存段里获取到一个新的slot
        // 必须有一个slot, 来存放 共享内存段
        if (slot != null) {
          return slot;
        }
        // 没有可用的槽位，但是已经有线程正在申请新的共享内存，则等待
        if (loading) {
          finishedLoading.awaitUninterruptibly();
        } else {
          // 当前线程申请新的共享内存
          loading = true;
          lock.unlock();
          DfsClientShm shm;
          try {
            // 通过DataTransferProtocol申请新的共享内存
            shm = requestNewShm(clientName, peer); // TODO
            if (shm == null) {
              continue;
            }
            // 将已经申请的共享内存加入 domainSocketWatcher，监控状态
            domainSocketWatcher.add(peer.getDomainSocket(), shm);
            usedPeer.setValue(true);
          } finally {
            lock.lock();
            loading = false;
            // 通知其他所有等待的线程
            finishedLoading.signalAll();
          }
          if (shm.isDisconnected()) {
            // 如果刚刚申请的共享内存已经过期了，那么继续循环
          } else {
            // 将刚刚申请的共享内存加入 notFull队列中
            notFull.put(shm.getShmId(), shm);
          }
        }
      }
    }
    
    /**
     *  freeSlot ()方法用于释放一个Slot，与allocSlot()方法的实现很类似。
     *  它首先会从共享内存中释放这个槽位，然后根据共享内存的状态执行对应的操作，
     */
    void freeSlot(Slot slot) {
      DfsClientShm shm = (DfsClientShm)slot.getShm();
      // 从共享内存中释放槽位
      shm.unregisterSlot(slot.getSlotIdx());
      if (shm.isDisconnected()) {
        Preconditions.checkState(!full.containsKey(shm.getShmId()));
        Preconditions.checkState(!notFull.containsKey(shm.getShmId()));
        // 如果共享内存已经被标记为过期，并且共享内存中的所有槽位都已经被释放，则将共享内存释放
        if (shm.isEmpty()) {
          shm.free();
        }
      } else {
        ShmId shmId = shm.getShmId();
        // 由于共享内存释放了一个槽位，那么这个共享内内存就不再是full状态了，从full队列中移除
        full.remove(shmId);
        // 如果共享内存中的所有槽位都空闲，则调用shutDown() 方法将共享内存相关的 domainSocket关闭。这时 DomainSocketWatcher 会调用 DfsClientShm.handle() 方法清理共享内存段
        if (shm.isEmpty()) {
          notFull.remove(shmId);
          shutdown(shm);
        } else {
          // 否则，就直接将共享内存放入notFull队列中。
          notFull.put(shmId, shm);
        }
      }
    }
    
    /**
     * Unregister a shared memory segment.
     *
     * Once a segment is unregistered, we will not allocate any more slots inside that segment.
     *
     * The DomainSocketWatcher calls this while holding the DomainSocketWatcher lock.
     *
     * @param shmId         The ID of the shared memory segment to unregister.
     */
    void unregisterShm(ShmId shmId) {
      lock.lock();
      try {
        full.remove(shmId);
        notFull.remove(shmId);
      } finally {
        lock.unlock();
      }
    }

    @Override
    public String toString() {
      return String.format("EndpointShmManager(%s, parent=%s)",
          datanode, DfsClientShmManager.this);
    }

    PerDatanodeVisitorInfo getVisitorInfo() {
      return new PerDatanodeVisitorInfo(full, notFull, disabled);
    }

    final void shutdown(DfsClientShm shm) {
      try {
        shm.getPeer().getDomainSocket().shutdown();
      } catch (IOException e) {
        LOG.warn(this + ": error shutting down shm: got IOException calling " +
            "shutdown(SHUT_RDWR)", e);
      }
    }
  }

  private boolean closed = false;

  private final ReentrantLock lock = new ReentrantLock();

  /**
   * A condition variable which is signalled when we finish loading a segment
   * from the Datanode.
   */
  private final Condition finishedLoading = lock.newCondition();

  /**
   * Information about each Datanode.
   */
  private final HashMap<DatanodeInfo, EndpointShmManager> datanodes =
      new HashMap<DatanodeInfo, EndpointShmManager>(1);
  
  /**
   * The DomainSocketWatcher which keeps track of the UNIX domain socket
   * associated with each shared memory segment.
   *
   * Note: because the DomainSocketWatcher makes callbacks into this
   * DfsClientShmManager object, you must MUST NOT attempt to take the
   * DomainSocketWatcher lock while holding the DfsClientShmManager lock,
   * or else deadlock might result.   This means that most DomainSocketWatcher
   * methods are off-limits unless you release the manager lock first.
   */
  private final DomainSocketWatcher domainSocketWatcher;
  
  DfsClientShmManager(int interruptCheckPeriodMs) throws IOException {
    this.domainSocketWatcher = new DomainSocketWatcher(interruptCheckPeriodMs);
  }

  public Slot allocSlot(DatanodeInfo datanode, DomainPeer peer, MutableBoolean usedPeer, ExtendedBlockId blockId, String clientName) throws IOException {
    lock.lock();
    try {
      if (closed) {
        LOG.trace(this + ": the DfsClientShmManager isclosed.");
        return null;
      }
      EndpointShmManager shmManager = datanodes.get(datanode);
      if (shmManager == null) {
        shmManager = new EndpointShmManager(datanode);
        datanodes.put(datanode, shmManager);
      }
      return shmManager.allocSlot(peer, usedPeer, clientName, blockId); // TODO
    } finally {
      lock.unlock();
    }
  }
  
  public void freeSlot(Slot slot) {
    lock.lock();
    try {
      DfsClientShm shm = (DfsClientShm)slot.getShm();
      shm.getEndpointShmManager().freeSlot(slot);
    } finally {
      lock.unlock();
    }
  }

  @VisibleForTesting
  public static class PerDatanodeVisitorInfo {
    public final TreeMap<ShmId, DfsClientShm> full;
    public final TreeMap<ShmId, DfsClientShm> notFull;
    public final boolean disabled;

    PerDatanodeVisitorInfo(TreeMap<ShmId, DfsClientShm> full,
        TreeMap<ShmId, DfsClientShm> notFull, boolean disabled) {
      this.full = full;
      this.notFull = notFull;
      this.disabled = disabled;
    }
  }

  @VisibleForTesting
  public interface Visitor {
    void visit(HashMap<DatanodeInfo, PerDatanodeVisitorInfo> info)
        throws IOException;
  }

  @VisibleForTesting
  public void visit(Visitor visitor) throws IOException {
    lock.lock();
    try {
      HashMap<DatanodeInfo, PerDatanodeVisitorInfo> info = 
          new HashMap<DatanodeInfo, PerDatanodeVisitorInfo>();
      for (Entry<DatanodeInfo, EndpointShmManager> entry :
            datanodes.entrySet()) {
        info.put(entry.getKey(), entry.getValue().getVisitorInfo());
      }
      visitor.visit(info);
    } finally {
      lock.unlock();
    }
  }

  /**
   * Close the DfsClientShmManager.
   */
  @Override
  public void close() throws IOException {
    lock.lock();
    try {
      if (closed) return;
      closed = true;
    } finally {
      lock.unlock();
    }
    // When closed, the domainSocketWatcher will issue callbacks that mark
    // all the outstanding DfsClientShm segments as stale.
    IOUtils.cleanup(LOG, domainSocketWatcher);
  }


  @Override
  public String toString() {
    return String.format("ShortCircuitShmManager(%08x)",
        System.identityHashCode(this));
  }

  @VisibleForTesting
  public DomainSocketWatcher getDomainSocketWatcher() {
    return domainSocketWatcher;
  }
}
